package org.example.controller;

import lombok.RequiredArgsConstructor;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.Media;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.ai.ollama.api.OllamaModel;
import org.springframework.core.io.ClassPathResource;
import org.springframework.util.MimeTypeUtils;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.util.List;

@RequiredArgsConstructor
@RestController
public class ApiController {

    private final ChatClient chatClient;

    private final OllamaChatModel chatModel;

    @GetMapping("/ai/chat")
    public String chat(@RequestParam("message") String message) {
        return chatClient.prompt().user(message).call().content();
    }

    @GetMapping("/ai/media")
    public String media() {
        ClassPathResource classPathResource = new ClassPathResource("test.jpg");
        Message message = new UserMessage(
                "请用中文描述下这张图片里有什么东西？",
                List.of(new Media(MimeTypeUtils.IMAGE_JPEG, classPathResource))
        );
        return chatModel.call(new Prompt(List.of(message), ChatOptions.builder().model(OllamaModel.LLAVA.getName()).build())).getResult().getOutput().getText();
    }
}
